"""
Paraphrase Multilingual MiniLM L12 V2 推理引擎
"""
import torch
import logging
import numpy as np
from typing import List, Dict, Optional, Union
from sentence_transformers.util import cos_sim

logger = logging.getLogger(__name__)


class ParaphraseMiniLMInference:
    """Paraphrase MiniLM推理引擎"""
    
    def __init__(self, model, device: str):
        self.model = model
        self.device = device
    
    def encode(
        self,
        texts: Union[str, List[str]],
        normalize_embeddings: bool = False,
        batch_size: int = 32,
        **kwargs
    ) -> np.ndarray:
        """
        编码文本为向量
        
        Args:
            texts: 文本或文本列表
            normalize_embeddings: 是否归一化向量
            batch_size: 批处理大小
            **kwargs: 其他参数
        
        Returns:
            嵌入向量数组
        """
        try:
            if isinstance(texts, str):
                texts = [texts]
            
            embeddings = self.model.encode(
                texts,
                normalize_embeddings=normalize_embeddings,
                batch_size=batch_size,
                convert_to_numpy=True,
                **kwargs
            )
            
            return embeddings
            
        except Exception as e:
            logger.error(f"编码失败: {str(e)}")
            raise
    
    def similarity(
        self,
        texts: List[str],
        normalize: bool = True
    ) -> float:
        """
        计算两个文本的相似度
        
        Args:
            texts: 两个文本的列表
            normalize: 是否归一化
        
        Returns:
            相似度分数 (0-1)
        """
        try:
            if len(texts) != 2:
                raise ValueError("需要提供两个文本")
            
            embeddings = self.encode(texts, normalize_embeddings=normalize)
            
            # 计算余弦相似度
            similarity = cos_sim(embeddings[0], embeddings[1]).item()
            
            return float(similarity)
            
        except Exception as e:
            logger.error(f"相似度计算失败: {str(e)}")
            raise

    def compute_similarities(
        self,
        source_sentence: str,
        sentences: List[str],
        normalize: bool = True
    ) -> List[float]:
        """
        计算源句子与多个目标句子的相似度
        
        Args:
            source_sentence: 源句子
            sentences: 目标句子列表
            normalize: 是否归一化
        
        Returns:
            相似度分数列表
        """
        try:
            # 编码源句子和目标句子
            source_embedding = self.encode([source_sentence], normalize_embeddings=normalize)[0]
            target_embeddings = self.encode(sentences, normalize_embeddings=normalize)
            
            # 计算相似度
            similarities = cos_sim(source_embedding, target_embeddings)[0]
            
            return [float(score) for score in similarities.tolist()]
            
        except Exception as e:
            logger.error(f"批量相似度计算失败: {str(e)}")
            raise
    
    def search(
        self,
        query: str,
        documents: List[str],
        top_k: int = 5,
        normalize: bool = True
    ) -> List[Dict]:
        """
        语义搜索
        
        Args:
            query: 查询文本
            documents: 文档列表
            top_k: 返回前k个结果
            normalize: 是否归一化
        
        Returns:
            搜索结果列表
        """
        try:
            # 编码查询和文档
            query_embedding = self.encode([query], normalize_embeddings=normalize)[0]
            doc_embeddings = self.encode(documents, normalize_embeddings=normalize)
            
            # 计算相似度
            similarities = cos_sim(query_embedding, doc_embeddings)[0]
            
            # 获取top_k结果
            top_indices = torch.topk(similarities, k=min(top_k, len(documents))).indices.tolist()
            
            results = []
            for idx in top_indices:
                results.append({
                    "document": documents[idx],
                    "score": float(similarities[idx].item()),
                    "index": idx
                })
            
            # 按分数降序排序
            results.sort(key=lambda x: x["score"], reverse=True)
            
            return results
            
        except Exception as e:
            logger.error(f"搜索失败: {str(e)}")
            raise


